# -*- coding: utf-8 -*-
"""
create by shenjun 2017-12-12
================================
朴素贝叶斯算法实现棋类AI算法
适用：
1. 能用特征表（或词汇表）表述且只用来描述棋盘上的格子
2. 能且只能用胜利、失败[、平手]表述分类
"""
from numpy import *
import pickle

WIN = 1
LOSS = -1
Draw = 0


def set2Vec(vocabList, inputSet):
    returnVec = [0] * len(vocabList)
    for word in inputSet:
        if word in vocabList:
            returnVec[vocabList.index(word)] += 1
        else:
            print "the word: %s is not in my Vocabulary!" % word
    return returnVec


def trainNB(trainMatrix, trainResult):
    numSimples = len(trainMatrix)
    numVocs = len(trainMatrix[0])
    pWin = trainResult.count(WIN) / float(numSimples)
    pLoss = trainResult.count(LOSS) / float(numSimples)
    pDraw = trainResult.count(Draw) / float(numSimples)
    pWinNum = ones(numVocs)
    pLossNum = ones(numVocs)
    pDrawNum = ones(numVocs)
    pWinDenom = 2
    pLossDenom = 2
    pDrawDenom = 2
    for i in range(numSimples):
        if trainResult[i] == WIN:
            pWinNum += trainMatrix[i]
            pWinDenom += sum(trainMatrix[i])
        elif trainResult[i] == LOSS:
            pLossNum += trainMatrix[i]
            pLossDenom += sum(trainMatrix[i])
        elif trainResult[i] == Draw:
            pDrawNum += trainMatrix[i]
            pDrawDenom += sum(trainMatrix[i])
    pWinVect = log(pWinNum / pWinDenom)
    pLossVect = log(pLossNum / pLossDenom)
    pDrawVect = log(pDrawNum / pDrawDenom)
    return [pWinVect, pDrawVect, pLossVect], [pWin, pDraw, pLoss]


def initSimple():
    return [['NA'] * 10], [0]


def classifyNB(vocabList, inputSet, VectArry, pClassArr):
    inputVec = set2Vec(vocabList, inputSet)
    pWinVec, pDrawVec, pLossVec = VectArry
    pExpWin, pExpDraw, pExpLoss = pClassArr
    pWin = sum(inputVec * pWinVec) + log(pExpWin)
    pDraw = sum(inputVec * pDrawVec) + log(pExpDraw)
    pLoss = sum(inputVec * pLossVec) + log(pExpLoss)
    pWinner = max(pWin, pDraw, pLoss)
    if pWinner == pWin:
        return WIN, pWin, pLoss
    elif pWinner == pDraw:
        return Draw, pWin, pLoss
    elif pWinner == pLoss:
        return LOSS, pWin, pLoss


def storeTrainRst(data, fn):
    f = open(fn, 'w')
    pickle.dump(data, f)
    f.close()


def loadTrainRst(fn):
    f = open(fn, 'r')
    data = pickle.load(f)
    f.close()
    return data


def Decision(aiSets, playerSets, valiableVocLst, vocabList, savedTrainRstFn):
    """
    约定：aiSets和playerSet各元素数组的最后1个子元素为该组元素尝试选择的特征值
    :param aiSets:
    :param playerSets:
    :param valiableVocLst:
    :param vocabList:
    :param savedTrainRstFn:
    :return:
    """
    savedTrainRst = loadTrainRst(savedTrainRstFn)
    VectArry = savedTrainRst[0]
    pClassArr = savedTrainRst[1]
    valiableVocBuff = []
    valiableVocBuff.extend(valiableVocLst)
    prior = None
    inferior = None
    bestStep = None
    noChoiseStep = None
    analyzedAIRst = analyze(aiSets, vocabList, VectArry, pClassArr)
    analyzedPlayerRst = analyze(playerSets, vocabList, VectArry, pClassArr)
    for vocab in analyzedAIRst:
        rst, pWin, pLoss = analyzedAIRst[vocab]
        if rst == WIN:
            if pWin > prior:
                prior = pWin
                bestStep = vocab
        # elif prior is None:
        #     if pWin > inferior:
        #         inferior = pWin
        #         bestStep = vocab
        elif rst == LOSS:
            valiableVocBuff.remove(vocab)
            if inferior is None or pLoss < inferior:
                inferior = pLoss
                noChoiseStep = vocab
    pPrior = None
    pBestStep = None
    for vocab in analyzedPlayerRst:
        rst, pWin, pLoss = analyzedPlayerRst[vocab]
        if rst == WIN:
            if pWin > pPrior:
                pPrior = pWin
                pBestStep = vocab
    # if bestStep is None:
    #     return valiableVocLst[int(random.uniform(0, len(valiableVocLst)))]
    # elif prior is None and inferior == log(0):
    #     return valiableVocLst[int(random.uniform(0, len(valiableVocLst)))]
    # else:
    #     return bestStep
    aiLog = open('ai.log', 'a')
    aiLogFmt = "choise[%s]\tai:pWin[%s]winStep[%s]pLoss[%s]worseStep[%s]\tplayer:pWin[%s]Step[%s]\n"
    if bestStep is None:
        if pBestStep is not None:
            aiLog.write(aiLogFmt
                        % (pBestStep, str(prior), str(bestStep), str(inferior), str(noChoiseStep), str(pPrior)
                           , str(pBestStep)))
            aiLog.close()
            return pBestStep
        elif len(valiableVocBuff) == 0 and noChoiseStep is None:
            stp = valiableVocLst[int(random.uniform(0, len(valiableVocLst)))]
            aiLog.write(aiLogFmt
                        % (stp, str(prior), str(bestStep), str(inferior), str(noChoiseStep), str(pPrior)
                           , str(pBestStep)))
            aiLog.close()
            return stp
        elif len(valiableVocBuff) > 0:
            stp = valiableVocBuff[int(random.uniform(0, len(valiableVocBuff)))]
            aiLog.write(aiLogFmt
                        % (stp, str(prior), str(bestStep), str(inferior), str(noChoiseStep), str(pPrior)
                           , str(pBestStep)))
            aiLog.close()
            return stp
        else:
            aiLog.write(aiLogFmt
                        % (noChoiseStep, str(prior), str(bestStep), str(inferior), str(noChoiseStep), str(pPrior)
                           , str(pBestStep)))
            aiLog.close()
            return noChoiseStep
    else:
        if pBestStep is not None:
            if pPrior > prior:
                aiLog.write(aiLogFmt
                            % (pBestStep, str(prior), str(bestStep), str(inferior), str(noChoiseStep), str(pPrior)
                               , str(pBestStep)))
                aiLog.close()
                return pBestStep
        aiLog.write(aiLogFmt
                    % (bestStep, str(prior), str(bestStep), str(inferior), str(noChoiseStep), str(pPrior)
                       , str(pBestStep)))
        aiLog.close()
        return bestStep


def analyze(inputSets, vocabList, VectArry, pClassArr):
    analyzedRst = {}
    for inputSet in inputSets:
        vocab = inputSet[-1]
        rst, pWin, pLoss = classifyNB(vocabList, inputSet[:-1], VectArry, pClassArr)
        analyzedRst[vocab] = [rst, pWin, pLoss]
    return analyzedRst


if __name__ == '__main__':
    VocLst = ['pA1', 'pA2', 'pA3',
              'pB1', 'pB2', 'pB3',
              'pC1', 'pC2', 'pC3',
              'oA1', 'oA2', 'oA3',
              'oB1', 'oB2', 'oB3',
              'oC1', 'oC2', 'oC3',
              'NA']
    data, cata = initSimple()

    VectArry, pClassArr = trainNB([set2Vec(VocLst, data[0])], cata)
    print VectArry
    print pClassArr
    storeTrainRst([VectArry, pClassArr], 'db/ai.db')
    obj = loadTrainRst('db/ai.db')
    print obj[0]
    print obj[1]
